#include <stdlib.h>

#include "../hle.h"

#include "__SysMemUserForUser.h"
#include "__ThreadManForUser.h"
#include "../blkalloc.h"
#include "../elfreader.h"
#include "../fs/metafs.h"

#include "__ModuleMgrForUser.h"

#define PSPLOADADDRESS 0x08800000

typedef struct
{
    /* 0, 0, 1, 1 ? */
    u16 moduleAttrs; /* 0x0000 User Mode, 0x1000 Kernel Mode */
    u16 moduleVersion;
    /* 28 bytes of module name, packed with 0's. */
    char name[28];
    u32 gp;           /* ptr to MIPS GOT data  (global offset table) */
    u32 libent;       /* ptr to .lib.ent section */
    u32 libentend;    /* ptr to end of .lib.ent section */
    u32 libstub;      /* ptr to .lib.stub section */
    u32 libstubend;   /* ptr to end of .lib.stub section */
} PspModuleInfo;

typedef struct
{
    u32 name; /* stub's name (module name) address */
    u16 version;
    u16 flags;
    u16 size;
    u16 imports; /* number of imports (syscalls) */
    u32 nid; /* address of the NIDs (the address pointing to where the executable's segments wrote the stub NIDs) */
    u32 text; /* address of where the syscall instructions should be written */
} Stub;

typedef struct
{
    u32 name; /* ent's name (module name) address */
    u16 version;
    u16 flags;
    u8 size;
    u8 vcount;
    u16 fcount;
    u32 resident;
} Ent;

SceUID moduleId;

void __KernelLoadStubs(u32 stubaddr, s32 nstub)
{
    s32 i, j, syscall, nid;
    Stub stubinfo;
    _log(INF, HLE, "%d stubs at 0x%08x", nstub, stubaddr);

    /* Reads all stubs (1 stub = all the definitions of one module's syscalls) */
    for (i = 0; i < nstub; i++)
    {
        char *name;

        mem_readStruct(stubaddr, &stubinfo, sizeof(Stub));

        stubaddr += sizeof(Stub);

        name = mem_readStr(stubinfo.name);
        _log(INF, HLE, "Loading stub %d named %s from %08x to %08x", i, name, stubinfo.nid, stubinfo.text);

        /* Invalid in or out address! */
        if (!mem_isOk(stubinfo.nid) || !mem_isOk(stubinfo.text + 4)) {
            _log(WRN, HLE, "Wrong address for stub, skipping");
            mem_freeStr(name);
            continue;
        }

        /* Read all the imports (syscalls) of the current stub */
        for (j = 0; j < stubinfo.imports; j++)
        {
            nid = mem_read32(stubinfo.nid + j * 4); /* the NID of the current import */
            syscall = hle_getNewSyscall(name, nid); /* syscall number (index from the syscalls list of this emulator) */
            mem_write32(stubinfo.text + j * 8 + 4, (syscall << 6) + 0xc); /* write the syscall instruction to the out address */
        }
        mem_freeStr(name);
    }
}

void __KernelLoadEnts(u32 entaddr, u32 nent, char *modname)
{
    u32 i, j;
    Ent entinfo;
    _log(INF, HLE, "%d ents at 0x%08x", nent, entaddr);
    for (i = 0; i < nent; i++)
    {
        char *name;
        
        mem_readStruct(entaddr, &entinfo, sizeof(Ent));
        if (entinfo.size == 0) {
            _log(INF, HLE, "Skipping dummy ent");
            continue;
        }
        if (entinfo.size > 4)
            entaddr += entinfo.size * 4;
        else
            entaddr += sizeof(Ent);

        if (entinfo.name == 0)
            name = modname;
        else
            name = mem_readStr(entinfo.name);
        _log(INF, HLE, "Exporting ent %d named %s, %d funcs, %d vars, resident %08x", i, name, entinfo.fcount, entinfo.vcount, entinfo.resident);

        for (j = 0; j < entinfo.fcount; j++)
        {
            u32 nid = mem_read32(entinfo.resident + j * 4);
            u32 exportAddr = mem_read32(entinfo.resident + (entinfo.fcount + entinfo.vcount) * 4 + j * 4);
            hle_addExport(name, nid, exportAddr);
            _log(INF, HLE, "Exporting 0x%08x (%s) at 0x%08x", nid, hle_getNidName(nid), exportAddr);
        }

        if (name != modname)
            mem_freeStr(name);
    }
}

SceUID __KernelLoadModule(Module *m, PspModuleInfo *modinfo, u32 modaddr)
{
    moduleId = ko_init(m, KO_MODULE);
    mem_readStruct(modaddr, modinfo, sizeof(PspModuleInfo));
    m->gp_value = modinfo->gp;
    strncpy(m->name, modinfo->name, 28);
    _log(INF, HLE, "Module at %08x: name %s, gp 0x%08x, attr 0x%04x, version 0x%04x", modaddr,
          modinfo->name, modinfo->gp, modinfo->moduleAttrs, modinfo->moduleVersion);
    m->encrypted = 0;
    return moduleId;
}

SceUID __KernelLoadEncryptedModule(Module *m)
{
    moduleId = ko_init(m, KO_MODULE);
    _log(INF, HLE, "Loading encrypted module");
    m->encrypted = 1;
    return moduleId;
}

void __KernelLoadFakeSyscalls()
{
    u32 retThAddr = blkAlloc_alloc(&kernelMemory, 4, 0, "RetFromThread", 0);
    u32 retCbAddr = blkAlloc_alloc(&kernelMemory, 4, 0, "RetFromCallback", 0);

    mem_write32(retThAddr, 
                 (hle_getNewSyscall("__KernelReturnFromThread", NID_THRET) << 6) + 0xc);
    __KernelSetThRetAddr(retThAddr);

    mem_write32(retCbAddr, 
                 (hle_getNewSyscall("__KernelReturnFromCallback", NID_CBRET) << 6) + 0xc);
    __KernelSetCbRetAddr(retCbAddr);
}

SceUID __KernelLoadELFFromPtr(Module *m, const char *ptr, s32 loadAddress)
{
    PspModuleInfo modinfo;
    u32 modaddr, stubaddr, entaddr;
    u16 modattr;
    SceUID id;
    s32 moduleSect;
    ElfReader reader;

    if (strncmp(ptr, "~PSP", 4) == 0
     || strncmp(ptr, "~SCE", 4) == 0)
        return __KernelLoadEncryptedModule(m);

    if (strncmp(ptr, "\x7F""ELF", 4) != 0) {
        _log(ERR, HLE, "File is not an ELF!");
        return -1;
    }

    /* Load the ELF file */
    reader = elfr_init((void*)ptr);

    /* Load module */
    moduleSect = elfr_getSectByName(&reader, ".rodata.sceModuleInfo");
    if (moduleSect != -1)
        modattr = *(u16*)(ptr + reader.sections[moduleSect].sh_offset);
    else
        modattr = *(u16*)(ptr + reader.segments[0].p_paddr - reader.segments[0].p_offset);

    if (loadAddress == -1)
    {
        if ((modattr & 0x1000) == 0)
        {
            _log(INF, HLE, "Loading user module");
            loadAddress = blkAlloc_alloc(&userMemory, reader.maxVaddr, 0, "TMP", 0);
            if (loadAddress)
                blkAlloc_free(&userMemory, loadAddress);
            else
                _log(ERR, HLE, "Can't load user module of size 0x%08x!", reader.maxVaddr);
        }
        else
        {
            _log(INF, HLE, "Loading kernel module");
            loadAddress = blkAlloc_alloc(&kernelMemory, reader.maxVaddr, 0, "TMP", 0);
            if (loadAddress)
                blkAlloc_free(&kernelMemory, loadAddress);
            else
                _log(ERR, HLE, "Can't load kernel module of size 0x%08x!", reader.maxVaddr);
        }
    }

    _log(INF, HLE, "Loading ELF at address 0x%08x", loadAddress);
    elfr_load(&reader, loadAddress);

    /* Find the entry address of the module */
    m->entry_addr = (reader.header.e_entry != 0xffffffff) ? (reader.vaddr + reader.header.e_entry) : 0xffffffff;
    if (!mem_isOk(m->entry_addr))
        _log(ERR, HLE, "Invalid module entry address!");

    _log(INF, HLE, "Module has entry address %08x", m->entry_addr);

    /* Load module */
    if (moduleSect != -1)
        modaddr = reader.vaddr + reader.sections[moduleSect].sh_addr;
    else
        modaddr = reader.vaddr + reader.segments[0].p_paddr - reader.segments[0].p_offset;
    if (!mem_isOk(modaddr)) {
        _log(ERR, HLE, "Invalid module address 0x%08x!", modaddr);
        return -1;
    }

    id = __KernelLoadModule(m, &modinfo, modaddr);

    /* Load stubs (imports) */
    stubaddr = modinfo.libstub;
    if (!mem_isOk(stubaddr))
        _log(ERR, HLE, "Invalid stub address 0x%08x!", stubaddr);
    else {
        u32 nstub = (modinfo.libstubend - modinfo.libstub) / sizeof(Stub);
        __KernelLoadStubs(stubaddr, nstub);
    }

    /* Load ents (exports) */
    entaddr = modinfo.libent;
    if (!mem_isOk(entaddr))
        _log(ERR, HLE, "Invalid ent address!");
    else {
        u32 nent = (modinfo.libentend - modinfo.libent) / sizeof(Ent);
        __KernelLoadEnts(entaddr, nent, m->name);
    }

    return id;
}

void __KernelStartModule(SceUID uid, s32 args, const char *argp, SceKernelSMOption *options)
{
    __KernelSetupRootThread(uid, args, argp, options->priority, options->stacksize, options->attribute);
    /* TODO: if current thread, put it in wait state, waiting for the new thread */
}

u32 __KernelGetModuleGP(SceUID module)
{
    Module *m = ko_get(module);
    if (m)
        return m->gp_value;
    return 0;
}

void __KernelSetDefaultSMOption(SceKernelSMOption *option)
{
    option->size = sizeof(SceKernelSMOption);
    option->attribute = 0;
    option->mpidstack = 2;
    option->priority = 32;
    option->stacksize = 0x4000;
}

SceUID __KernelLoadELF(Module *m, const char *filename, s32 loadAddress)
{
    FileInfo info;
    u32 size;
    u32 handle;
    char *temp;
    SceUID id;

    /* Load file */
    info = mfs_getFileInfo(filename);
    size = info.size;

    handle = mfs_openFile(filename, FILEACCESS_READ);
   
    if (!size) {
        _log(ERR, HLE, "Problem while loading file %s!", filename);
        return -1;
    }
   
    temp = malloc(sizeof(char) * size);
    mfs_readFile(handle, temp, (size_t)size);
   
    /* Load the module & set PC */
    m = malloc(sizeof(Module));
    id = __KernelLoadELFFromPtr(m, temp, loadAddress);
   
    /* Close the file */
    free(temp);
    mfs_closeFile(handle);

    return id;
}

SceUID __KernelLoadELFFromHandle(Module *m, u32 handle, s32 loadAddress)
{
    u32 size = 0x1000000; /* TODO: VERY hacky */
    char *temp;
    SceUID id;
   
    if (!size) {
        _log(ERR, HLE, "Problem while loading file handle %d!", handle);
        return -1;
    }
   
    temp = malloc(sizeof(char) * size);
    mfs_readFile(handle, temp, (size_t)size);
   
    /* Load the module */
    m = malloc(sizeof(Module));
    id = __KernelLoadELFFromPtr(m, temp, loadAddress);
   
    /* Close the file */
    free(temp);

    return id;
}

LoadResult __KernelLoadExec(const char *filename)
{
    Module *m = malloc(sizeof(Module));
    SceKernelSMOption option;
    SceUID id = __KernelLoadELF(m, filename, PSPLOADADDRESS);
    if (id == -1)
        return LOAD_FAIL;

    /* Load the fake syscalls */
    __KernelLoadFakeSyscalls();

    /* Set module options & run it */
    __KernelSetDefaultSMOption(&option);
    __KernelStartModule(id, strlen(filename) + 1, filename, &option);

    return LOAD_OK;
}

LoadResult __KernelLoadExecFromPtr(const char *ptr, const char *filename)
{
    Module *m = malloc(sizeof(Module));
    SceKernelSMOption option;
    SceUID id = __KernelLoadELFFromPtr(m, ptr, PSPLOADADDRESS);
    if (id == -1)
        return LOAD_FAIL;

    /* Load the fake syscalls */
    __KernelLoadFakeSyscalls();

    __KernelSetDefaultSMOption(&option);
    __KernelStartModule(id, strlen(filename) + 1, filename, &option);

    return LOAD_OK;
}

SceUID __KernelGetCurModuleId()
{
    return moduleId;
}

